import numpy as np

forwardRate = 0.1   # 朝着正前方移动的概率
rightRate = 0.45     # 垂直正前方向左移动的概率
leftRate = 0.45      # 垂直正前方向右移动的概率

class Env():
    def __init__(self, name):
        self.Name = name
        self.N = 11
        self.A = np.arange(4)  # 0-右，1-下，2-左，3-上
        self.X = np.arange(self.N)   # 状态集合
        self.P = self.makeP()
        self.R = self.makeR()
        self.gamma = 0.9 # 折扣
        self.StartState = 0
        self.EndState = [6, 10]

    def makeP(self):
        # 计算得到转移矩阵
        p = np.zeros((self.N, 4, self.N))  # 在x状态下，采取了y行动，到达z(x')的概率
        for i in range(self.N):  # 根据每个点来遍历信息
            if i == 6 or i == 10:  # 终止点
                continue
            for move in self.A:
                # 判断经过move操作时候到达哪个状态了
                current = None
                k = 1.0  # 总的概率
                if move == 3: # 向上走
                    if i == 1 or i >= 7: current = -1
                    elif i >= 2 and i <= 4: current = i + 3
                    else: current = i + 4
                if move == 1: # 向下走
                    if i <= 3 or i == 8: current = -1
                    elif i >= 5 and i <= 7: current = i - 3
                    else: current = i - 4
                if move == 2:  # 向左走
                    if i in [0, 4, 5, 7]:
                        current = -1
                    else:
                        current = i - 1
                if move == 0: # 向右走
                    if i in [3, 4, 6, 10]:
                        current = -1
                    else:
                        current = i + 1
                # current=-1表示碰到了墙壁
                if current != -1:
                    p[i][move][current] += forwardRate  # 到达前进的点的概率
                    k -= forwardRate
                # 还有到达垂直与前进方向的点的概率
                if move == 3 or move == 1:  # 上下走
                    if i not in [0, 4, 5, 7]:  # 向左走不会撞墙
                        p[i][move][i - 1] += leftRate
                        k -= leftRate
                    if i not in [3, 4, 6, 10]:  # 向右走不会撞墙
                        p[i][move][i + 1] += rightRate
                        k -= rightRate
                else:
                    if i != 8 and i > 3: # 向下走不会撞墙
                        target = 0
                        if i >= 5 and i <= 7:
                            target = i - 3
                        else:
                            target = i - 4
                        p[i][move][target] += rightRate
                        k -= leftRate   # 如果左右概率不同的话此处也是一个变量，因为前进的方向不同，左右也会相反
                    if i != 1 and i < 7: # 向上走不会撞墙
                        target = 0
                        if i >= 2 and i <= 4:
                            target = i + 3
                        else:
                            target = i + 4
                        p[i][move][target] += rightRate
                        k -= rightRate
                # 判断停在原地的概率
                p[i][move][i] += k
        return p

    def makeR(self):
        R = np.zeros(self.N)
        for i in range(self.N):
            # 报酬在这里修改
            if i == 6: R[i] = -1
            elif i == 10: R[i] = 1
            else: R[i] = -0.04
        return R

# 价值迭代算法
def ValueIter(E):
    U = np.zeros(E.N)
    U_ = np.zeros(E.N)
    delta = 1
    while delta > 0.0001:
        U = np.copy(U_)
        U_ = E.R + E.gamma * np.max(np.dot(E.P[:, :, :], U), axis=1)
        delta = np.max(np.abs(U-U_))
    Pai = np.argmax(np.dot(E.P[:, :, :], U), axis=1)
    return U, Pai

# 根据策略的收益来评估策略
def Eval(E, Pai):
    U = np.zeros(E.N)
    k = 20
    for k in range(k):
        for i in range(E.N):
            U[i] = E.R[i] + E.gamma * np.dot(E.P[i, Pai[i], :], U)
    return U

# 策略迭代算法
def PolicyIter(E):
    Pai = np.zeros(E.N, dtype=int)   # 初始的策略
    change = True
    while change:
        U = Eval(E, Pai)
        change = False
        for x in E.X:
            if np.max(np.dot(E.P[x, :, :], U)) > np.dot(E.P[x, Pai[x], :], U) + 1E-5:
                Pai[x] = np.argmax(np.dot(E.P[x, :, :], U))
                change = True
    return U, Pai

def mytest(E):
    for i in range(11):
        for j in range(4):
            for k in range(11):
                print('编号为{}的点经过操作{}到达{}的概率为：{}'.format(i, j, k, E.P[i][j][k]))

if __name__ == '__main__':
    E = Env(name="hlt is a good man")
    # mytest(E)
    # for i in E.A:
    #     print(i)
    # 采用价值迭代算法
    # U, Pai = ValueIter(E)

    U, Pai = PolicyIter(E)

    print("策略迭代算法得到的价值函数：", end=" ")
    print(U)
    print("----------------------------------------------")
    # print("采用策略迭代算法状态转移：", end=" ")
    print("采用价值迭代算法状态转移：", end=" ")
    print(Pai)

    print("------------------进行转化----------------------")


    for i in range(E.N):
        s=""
        if(Pai[i] == 0):
            s="向右走"
        elif(Pai[i] == 1):
            s="向下走"
        elif(Pai[i] == 2):
            s="向左走"
        elif(Pai[i] == 3):
            s="向上走"

        print("编号为"+str(i)+"的格子的策略为:"+s)


